#!/bin/bash
"""模型下载脚本

用于下载项目所需的预训练模型和模型检查点。
支持下载DeepSeek模型、BERT模型和RNN模型等。
"""

# 配置参数
MODEL_SAVE_DIR="${MODEL_SAVE_DIR:-./model_saves}"
DEEPSEEK_MODEL_URL="${DEEPSEEK_MODEL_URL:-https://example.com/models/deepseek-7b.zip}"
BERT_MODEL_URL="${BERT_MODEL_URL:-https://example.com/models/bert-crf.zip}"
RNN_MODEL_URL="${RNN_MODEL_URL:-https://example.com/models/rnn-sentiment.zip}"
TMP_DIR="${TMP_DIR:-./tmp}"

# 日志函数
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# 确保目录存在
ensure_dirs() {
    if [ ! -d "$MODEL_SAVE_DIR" ]; then
        log "创建模型保存目录: $MODEL_SAVE_DIR"
        mkdir -p "$MODEL_SAVE_DIR"
    fi
    
    if [ ! -d "$TMP_DIR" ]; then
        log "创建临时目录: $TMP_DIR"
        mkdir -p "$TMP_DIR"
    fi
}

# 下载模型文件
 download_model() {
    local model_name=$1
    local model_url=$2
    local model_dir="$MODEL_SAVE_DIR/$model_name"
    local tmp_file="$TMP_DIR/${model_name}.zip"
    
    # 检查模型是否已存在
    if [ -d "$model_dir" ]; then
        log "模型已存在: $model_name，跳过下载"
        return 0
    fi
    
    log "开始下载模型: $model_name 从 $model_url"
    
    # 下载模型文件
    if command -v wget &> /dev/null; then
        wget -O "$tmp_file" "$model_url"
    elif command -v curl &> /dev/null; then
        curl -o "$tmp_file" "$model_url"
    else
        log "错误: 未找到wget或curl命令，无法下载模型"
        return 1
    fi
    
    if [ $? -ne 0 ]; then
        log "错误: 模型下载失败"
        rm -f "$tmp_file" 2>/dev/null
        return 1
    fi
    
    log "模型下载完成，开始解压..."
    
    # 创建模型目录
    mkdir -p "$model_dir"
    
    # 解压模型文件
    unzip -q "$tmp_file" -d "$model_dir"
    
    if [ $? -ne 0 ]; then
        log "错误: 模型解压失败"
        rm -rf "$model_dir" 2>/dev/null
        rm -f "$tmp_file" 2>/dev/null
        return 1
    fi
    
    # 清理临时文件
    rm -f "$tmp_file"
    
    log "模型安装完成: $model_dir"
    return 0
}

# 验证模型完整性
validate_model() {
    local model_name=$1
    local model_dir="$MODEL_SAVE_DIR/$model_name"
    
    # 这里可以添加模型验证逻辑，例如检查关键文件是否存在
    case $model_name in
        "deepseek-7b")
            required_files=("config.json" "pytorch_model.bin" "tokenizer.json")
            ;;
        "bert-crf")
            required_files=("config.json" "model.bin" "tokenizer.json")
            ;;
        "rnn-sentiment")
            required_files=("config.json" "model.bin" "vocab.txt")
            ;;
        *)
            log "警告: 没有针对模型 $model_name 的验证规则"
            return 0
            ;;
    esac
    
    local missing_files=()
    for file in "${required_files[@]}"; do
        if [ ! -f "$model_dir/$file" ]; then
            missing_files+=($file)
        fi
    done
    
    if [ ${#missing_files[@]} -ne 0 ]; then
        log "错误: 模型 $model_name 缺少必要文件: ${missing_files[*]}"
        return 1
    fi
    
    log "模型 $model_name 验证通过"
    return 0
}

# 创建版本文件
create_version_file() {
    local version_file="$MODEL_SAVE_DIR/version.json"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    cat > "$version_file" << EOF
{
    "models": {
        "deepseek-7b": {
            "url": "$DEEPSEEK_MODEL_URL",
            "download_time": "$timestamp"
        },
        "bert-crf": {
            "url": "$BERT_MODEL_URL",
            "download_time": "$timestamp"
        },
        "rnn-sentiment": {
            "url": "$RNN_MODEL_URL",
            "download_time": "$timestamp"
        }
    },
    "last_updated": "$timestamp"
}
EOF
    
    log "创建版本文件: $version_file"
}

# 显示帮助信息
show_help() {
    echo "用法: $0 [选项]"
    echo "模型下载脚本，用于下载项目所需的预训练模型"
    echo ""
    echo "选项:"
    echo "  -m, --model <模型名>   要下载的模型: deepseek, bert, rnn, all (默认: all)"
    echo "  -d, --dir <目录>       模型保存目录 (默认: ./model_saves)"
    echo "  -h, --help             显示帮助信息"
    echo ""
    echo "环境变量:"
    echo "  MODEL_SAVE_DIR         同 --dir 选项"
    echo "  DEEPSEEK_MODEL_URL     DeepSeek模型下载URL"
    echo "  BERT_MODEL_URL         BERT模型下载URL"
    echo "  RNN_MODEL_URL          RNN模型下载URL"
}

# 解析命令行参数
parse_args() {
    local model_type="all"
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            -m|--model)
                model_type=$2
                shift 2
                ;;
            -d|--dir)
                MODEL_SAVE_DIR=$2
                shift 2
                ;;
            -h|--help)
                show_help
                exit 0
                ;;
            *)
                log "错误: 未知选项: $1"
                show_help
                exit 1
                ;;
        esac
    done
    
    echo "$model_type"
}

# 主函数
main() {
    log "===== 模型下载脚本启动 ====="
    
    # 解析命令行参数
    local model_type=$(parse_args "$@")
    
    # 确保目录存在
    ensure_dirs
    
    # 下载指定的模型
    local all_success=true
    
    case $model_type in
        "deepseek")
            download_model "deepseek-7b" "$DEEPSEEK_MODEL_URL"
            validate_model "deepseek-7b"
            all_success=$?
            ;;
        "bert")
            download_model "bert-crf" "$BERT_MODEL_URL"
            validate_model "bert-crf"
            all_success=$?
            ;;
        "rnn")
            download_model "rnn-sentiment" "$RNN_MODEL_URL"
            validate_model "rnn-sentiment"
            all_success=$?
            ;;
        "all")
            log "开始下载所有模型..."
            download_model "deepseek-7b" "$DEEPSEEK_MODEL_URL"
            validate_model "deepseek-7b"
            local deepseek_success=$?
            
            download_model "bert-crf" "$BERT_MODEL_URL"
            validate_model "bert-crf"
            local bert_success=$?
            
            download_model "rnn-sentiment" "$RNN_MODEL_URL"
            validate_model "rnn-sentiment"
            local rnn_success=$?
            
            all_success=$((deepseek_success && bert_success && rnn_success))
            ;;
        *)
            log "错误: 不支持的模型类型: $model_type"
            show_help
            exit 1
            ;;
    esac
    
    # 创建版本文件
    if [ $all_success -eq 0 ]; then
        create_version_file
    fi
    
    if [ $all_success -eq 0 ]; then
        log "===== 模型下载脚本执行成功 ====="
        exit 0
    else
        log "===== 模型下载脚本执行失败 ====="
        exit 1
    fi
}

# 执行主函数
main "$@"